#! /usr/bin/env python
# -*- coding: UTF-8 -*-
#===============================================================================#
#   DESCRIPTION:  
# 
#       OPTIONS:  ---
#  REQUIREMENTS:  python(of course)
#         NOTES:  write for liyan
#                 use this for special purpose
#        AUTHOR:  ybyygu (Wenping Guo) 
#       LICENCE:  GPL version 2 or upper
#       VERSION:  0.1
#       CREATED:  the last day of 2006 
#      REVISION:  ---
#===============================================================================#
import sys
import re

#===============================================================================#
#
#  Functions
#
#===============================================================================#

def getNeighboringAtomXYZ(eline, list):
    min = 5000.0
    t = eline.split()
    e = t[0]
    xx = float(t[1])
    yy = float(t[2])
    zz = float(t[3])
    for m in list:
        t = m.split()
        x = float(t[1])
        y = float(t[2])
        z = float(t[3])
        t = (xx - x)**2 + (yy - y)**2 + (zz -z)**2
        if t < min:
            min = t
            res = m
    return res

def getEquivalentAtom(m1):
    list = elexyz[m1.split()[0]]
    res = getNeighboringAtomXYZ(m1, list)
    return res

def removeBorderHAtoms():
    global mol1
    al = ''
    t = filter(lambda m: m[:2] == 'Al', mol1)
    if len(t) == 1:
        al = t[0]
        t, xx, yy, zz = al.split()
        xx = float(xx)
        yy = float(yy)
        zz = float(zz)
    if al:
        min = 50.0
        res = None
        for m in mol1:
            t, x, y, z = m.split()
            if t != 'H':
                continue
            x = float(x)
            y = float(y)
            z = float(z)
            t = (x - xx)**2 + (y - yy)**2 + (z - zz)**2
            if t < min:
                min = t
                res = m
    tl = filter(lambda m: m[0] != 'H', mol1)
    if al:
        tl.append(res)
    mol1 = tl

#===============================================================================#
#
#  Main
#
#===============================================================================#

xyz_file = sys.argv[1]
txt = open(xyz_file).read()
txt = txt.replace('\r', '')
txt = txt.strip()
mols = re.split('\n\d+\n', txt)
mol1 = mols[0].split('\n')
del mol1[0]
mol2 = mols[1].split('\n')
del mol1[0]
del mol2[0]

###
# the small one should be the first
#
if len(mol1) > len(mol2):
    mol1, mol2 = mol2, mol1

###
# make a xyz dictionary of different elements for further querying
#
elexyz = {}
for m in mol2:
    t = m.split()
    if elexyz.has_key(t[0]):
        elexyz[t[0]].append(m)
    else:
        elexyz[t[0]] = []
        elexyz[t[0]].append(m)

removeBorderHAtoms()
for m1 in mol1:
    m2 = getEquivalentAtom(m1)
    mol2[mol2.index(m2)] = m1

###
# Output
#
filename = xyz_file[:-4] + '-comoled.gjf'
fp = open(filename, 'w')
fp.writelines('# HF\n\n')
fp.writelines('generated by comol.py for liyan\n\n0 1\n')
fp.writelines('\n'.join(mol2))
fp.close()

print 'Liyan: Please check %s, and then open it in gaussian view.' % filename
